#include "ImgProcess.hpp"
#include "ImgProcess.cuh"
#include "raycast/raycast.h"
typedef float4 Point;
typedef uchar4 RGB;
typedef uint16_t Depth;
typedef float3 Normal;
typedef float3 Vertex;

__global__ void depth2vertex(Patch<Vertex> vertex, const Patch<Depth> depth, const Camera cam)
{
    const uint2 pixel = thr2pos2();
    // if (pixel.x >= depth.rows || pixel.y >= depth.cols)
    //     return;

    if (depth[pixel] > 0)
    {
        vertex[pixel] = cam.cam2world(pixel.x, pixel.y, depth[pixel] * 0.001f); // depth[pixel] * (rotate(invK, make_float3(pixel.x, pixel.y, 1.f)));
    }
    else
    {
        vertex[pixel] = make_float3(0, 0, 0);
    }
}

__global__ void vertex2normal(Patch<Normal> normal, const Patch<Vertex> vertex)
{
    const uint2 pixel = thr2pos2();
    // if (pixel.x >= vertex.rows || pixel.y >= vertex.cols)
    //     return;

    const float3 left = vertex[make_uint2(max(int(pixel.x) - 1, 0), pixel.y)];
    const float3 right = vertex[make_uint2(min(pixel.x + 1, vertex.rows - 1), pixel.y)];
    const float3 up = vertex[make_uint2(pixel.x, max(int(pixel.y) - 1, 0))];
    const float3 down = vertex[make_uint2(pixel.x, min(pixel.y + 1, vertex.cols - 1))];

    if (left.z == 0 || right.z == 0 || up.z == 0 || down.z == 0)
    {
        normal[pixel].x = CUDART_NAN_F;
        return;
    }

    const float3 dxv = right - left;
    const float3 dyv = down - up;
    normal[pixel] = normalized(cross(dyv, dxv)); // switched dx and dy to get factor -1
}

////////////////////////////////////////////////////////////////////////////////
// clamp
// - clamp the value v to be in the range [a, b]
////////////////////////////////////////////////////////////////////////////////

inline __device__ __host__ float clamp(float f, float a, float b)
{
    return fmaxf(a, fminf(f, b));
}
inline __device__ __host__ int clamp(int f, int a, int b)
{
    return max(a, min(f, b));
}
inline __device__ __host__ uint clamp(uint f, uint a, uint b)
{
    return max(a, min(f, b));
}

inline __device__ __host__ float2 clamp(float2 v, float a, float b)
{
    return make_float2(clamp(v.x, a, b), clamp(v.y, a, b));
}
inline __device__ __host__ float2 clamp(float2 v, float2 a, float2 b)
{
    return make_float2(clamp(v.x, a.x, b.x), clamp(v.y, a.y, b.y));
}
inline __device__ __host__ float3 clamp(float3 v, float a, float b)
{
    return make_float3(clamp(v.x, a, b), clamp(v.y, a, b), clamp(v.z, a, b));
}
inline __device__ __host__ float3 clamp(float3 v, float3 a, float3 b)
{
    return make_float3(clamp(v.x, a.x, b.x), clamp(v.y, a.y, b.y), clamp(v.z, a.z, b.z));
}
inline __device__ __host__ float4 clamp(float4 v, float a, float b)
{
    return make_float4(clamp(v.x, a, b), clamp(v.y, a, b), clamp(v.z, a, b), clamp(v.w, a, b));
}
inline __device__ __host__ float4 clamp(float4 v, float4 a, float4 b)
{
    return make_float4(clamp(v.x, a.x, b.x), clamp(v.y, a.y, b.y), clamp(v.z, a.z, b.z), clamp(v.w, a.w, b.w));
}

inline __device__ __host__ int2 clamp(int2 v, int a, int b)
{
    return make_int2(clamp(v.x, a, b), clamp(v.y, a, b));
}
inline __device__ __host__ int2 clamp(int2 v, int2 a, int2 b)
{
    return make_int2(clamp(v.x, a.x, b.x), clamp(v.y, a.y, b.y));
}
inline __device__ __host__ int3 clamp(int3 v, int a, int b)
{
    return make_int3(clamp(v.x, a, b), clamp(v.y, a, b), clamp(v.z, a, b));
}
inline __device__ __host__ int3 clamp(int3 v, int3 a, int3 b)
{
    return make_int3(clamp(v.x, a.x, b.x), clamp(v.y, a.y, b.y), clamp(v.z, a.z, b.z));
}
inline __device__ __host__ int4 clamp(int4 v, int a, int b)
{
    return make_int4(clamp(v.x, a, b), clamp(v.y, a, b), clamp(v.z, a, b), clamp(v.w, a, b));
}
inline __device__ __host__ int4 clamp(int4 v, int4 a, int4 b)
{
    return make_int4(clamp(v.x, a.x, b.x), clamp(v.y, a.y, b.y), clamp(v.z, a.z, b.z), clamp(v.w, a.w, b.w));
}

inline __device__ __host__ uint2 clamp(uint2 v, uint a, uint b)
{
    return make_uint2(clamp(v.x, a, b), clamp(v.y, a, b));
}
inline __device__ __host__ uint2 clamp(uint2 v, uint2 a, uint2 b)
{
    return make_uint2(clamp(v.x, a.x, b.x), clamp(v.y, a.y, b.y));
}
inline __device__ __host__ uint3 clamp(uint3 v, uint a, uint b)
{
    return make_uint3(clamp(v.x, a, b), clamp(v.y, a, b), clamp(v.z, a, b));
}
inline __device__ __host__ uint3 clamp(uint3 v, uint3 a, uint3 b)
{
    return make_uint3(clamp(v.x, a.x, b.x), clamp(v.y, a.y, b.y), clamp(v.z, a.z, b.z));
}
inline __device__ __host__ uint4 clamp(uint4 v, uint a, uint b)
{
    return make_uint4(clamp(v.x, a, b), clamp(v.y, a, b), clamp(v.z, a, b), clamp(v.w, a, b));
}
inline __device__ __host__ uint4 clamp(uint4 v, uint4 a, uint4 b)
{
    return make_uint4(clamp(v.x, a.x, b.x), clamp(v.y, a.y, b.y), clamp(v.z, a.z, b.z), clamp(v.w, a.w, b.w));
}

// inline __host__ __device__ float3 operator*(const float3 a, float b) const
// {
//     return make_float3(a.x * b, a.y * b, a.z * b);
// }

__global__ void renderLightKernel(Patch<RGB> out, const Patch<Vertex> vertex, const Patch<Normal> normal, const float3 light, const float3 ambient)
{
    const uint2 pixel = make_uint2(threadIdx.x, blockIdx.x);

    if (fabs(normal[pixel].z) < 1e-4)
    {
        const float3 bgr1 = make_float3(4.f / 255.f, 2.f / 255.f, 2.f / 255.f);
        const float3 bgr2 = make_float3(236.f / 255.f, 120.f / 255.f, 120.f / 255.f);

        const float w = static_cast<float>(pixel.x) / out.rows;
        auto color = bgr1 * (1 - w) + bgr2 * w;
        out[pixel] = make_uchar4(color.x * 255, color.y * 255, color.z * 255, 255);
    }

    else
    {
        const float3 diff = normalized(light - vertex[pixel]);
        const float dir = fmaxf(dot(normal[pixel], diff), 0.f);
        float3 col = clamp(make_float3(dir, dir, dir) + ambient, 0.f, 1.f);
        out[pixel] = make_uchar4(col.x * 255, col.y * 255, col.z * 255, 255);
    }
}

__global__ void render_image_kernel(const Patch<Depth> depth, const Patch<float3> normals,
                                    const Camera cam, Patch<RGB> dst)
{
    const uint2 pixel = make_uint2(threadIdx.x, blockIdx.x);
    uint x = pixel.x, y = pixel.y;

    float3 color;

    int d = depth[pixel];

    if (d == 0)
    {
        const float3 bgr1 = make_float3(4.f / 255.f, 2.f / 255.f, 2.f / 255.f);
        const float3 bgr2 = make_float3(236.f / 255.f, 120.f / 255.f, 120.f / 255.f);

        const float w = static_cast<float>(x) / dst.rows;
        color = bgr1 * (1 - w) + bgr2 * w;
    }
    else
    {
        const float3 P = cam.cam2world(x, y, d * 0.0002f);
        const float3 v4 = normals[pixel];

        const float3 N = make_float3(v4.x, v4.y, v4.z);

        const float Ka = 0.3f; // ambient coeff
        const float Kd = 0.5f; // diffuse coeff
        const float Ks = 0.2f; // specular coeff
        const float n = 20.f;  // specular power

        const float Ax = 1.f; // ambient color,  can be RGB
        const float Dx = 1.f; // diffuse color,  can be RGB
        const float Sx = 1.f; // specular color, can be RGB
        const float Lx = 1.f; // light color

        // Ix = Ax*Ka*Dx + Att*Lx [Kd*Dx*(N dot L) + Ks*Sx*(R dot V)^n]

        float3 L = normalized(make_float3(0, 0, 0) - P);
        float3 V = normalized(make_float3(0.f, 0.f, 0.f) - P);
        float3 R = normalized(make_float3(N.x * 2.0f, N.y * 2.0f, N.z * 2.0f) * dot(N, L) - L);

        float Ix = Ax * Ka * Dx + Lx * Kd * Dx * fmax(0.f, dot(N, L)) + Lx * Ks * Sx * __powf(fmax(0.f, dot(R, V)), n);
        color = make_float3(Ix, Ix, Ix);
    }

    RGB out;
    out.x = static_cast<unsigned char>(__saturatef(color.x) * 255.f);
    out.y = static_cast<unsigned char>(__saturatef(color.y) * 255.f);
    out.z = static_cast<unsigned char>(__saturatef(color.z) * 255.f);
    dst[pixel] = out;
}

__global__ void tangent_colors_kernel2(const Patch<Normal> normals, Patch<RGB> colors)
{
    const uint2 pixel = make_uint2(threadIdx.x, blockIdx.x);
    uint x = pixel.x, y = pixel.y;

    float3 n = normals[pixel];
#if 1
            unsigned char r = static_cast<unsigned char>(__saturatef((-n.x + 1.f)/2.f) * 255.f);
            unsigned char g = static_cast<unsigned char>(__saturatef((-n.y + 1.f)/2.f) * 255.f);
            unsigned char b = static_cast<unsigned char>(__saturatef((-n.z + 1.f)/2.f) * 255.f);
#else
    unsigned char r = static_cast<unsigned char>((5.f - n.x * 3.5f) * 25.5f);
    unsigned char g = static_cast<unsigned char>((5.f - n.y * 2.5f) * 25.5f);
    unsigned char b = static_cast<unsigned char>((5.f - n.z * 3.5f) * 25.5f);
#endif
    colors[pixel] = make_uchar4(b, g, r, 0);
}

void depth2img(Patch<Depth> depth, Camera &cam, Patch<Vertex> &vertex, Patch<Normal> &normal, Patch<RGB> out)
{

    vertex.create(depth.rows, depth.cols);
    normal.create(depth.rows, depth.cols);
    // out.creat(depth.rows, depth.cols);

    depth2vertex<<<depth.cols, depth.rows>>>(vertex, depth, cam);
    ck(cudaGetLastError());

    vertex2normal<<<depth.cols, depth.rows>>>(normal, vertex);
    ck(cudaGetLastError());

    const float3 light = make_float3(0, 2, 0);
    const float3 ambient = make_float3(0.3, 0.5, 0.2);
    renderLightKernel<<<depth.cols, depth.rows>>>(out, vertex, normal, light, ambient);
    // render_image_kernel<<<depth.cols, depth.rows>>>(depth, normal, cam, out);
    // tangent_colors_kernel<<<depth.cols, depth.rows>>>(normal, out);
    ck(cudaGetLastError());
}